// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <cantor/function/registry.h>
#include <turbo/utility/internal/status.h>

namespace cantor {

    turbo::Status FunctionRegistry::register_function(const std::shared_ptr<cantor::pb::catalog::FunctionDefine> &meta,
                                        const std::shared_ptr<FunctionCaller> &caller) {
        auto id = meta->id();
        auto dit = _function_defines.find(id);
        if (dit != _function_defines.end()) {
            return turbo::already_exists_error("function already exists");
        }

        auto sit = _function_sets.find(meta->func_name());
        if (sit == _function_sets.end()) {
            if (meta->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_SIMD ||meta->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_GPU) {
                return turbo::already_exists_error("function must have simple function version");
            }
            auto set = std::make_shared<FunctionSet>();
            set->func_name = meta->func_name();
            _function_sets[meta->func_name()] = set;
            sit = _function_sets.find(meta->func_name());
        }
        bool have_simpled = false;
        if (meta->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_SIMD ||meta->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_GPU) {
            for (auto &it : sit->second->functions) {
                if (it->feature() == cantor::pb::catalog::FunctionFeature::FUNC_FEATURE_NONE && it->base_id() == meta->base_id()) {
                    have_simpled = true;
                    break;
                }
            }
            if (!have_simpled) {
                return turbo::already_exists_error("function must have simple function version");
            }
        }
        sit->second->functions.push_back(meta);
        _function_defines[id] = meta;
        _function_callers[id] = caller;
        return turbo::OkStatus();
    }

    std::shared_ptr<FunctionSet> FunctionRegistry::find_function_set(const std::string &func_name) const {
        auto it = _function_sets.find(func_name);
        if (it == _function_sets.end()) {
            return nullptr;
        }
        return it->second;
    }


    std::shared_ptr<FunctionCaller> FunctionRegistry::find_caller(uint64_t id) const {
        auto it = _function_callers.find(id);
        if (it == _function_callers.end()) {
            return nullptr;
        }
        return it->second;
    }

    std::shared_ptr<cantor::pb::catalog::FunctionDefine> FunctionRegistry::find_function_define(uint64_t id) const {
        auto it = _function_defines.find(id);
        if (it == _function_defines.end()) {
            return nullptr;
        }
        return it->second;
    }
}  // namespace cantor
